#include "libmri.h"
#include <cublas.h>

#include <math.h>
#include <stdlib.h>
#include <iostream>

using std::cout;
using std::endl;
using std::cerr;
using std::flush;


void
calib3d_culanczos 
	(cplxf *&V, int& nv,
	 cplxf *_calib, int *csize, int *ksize, int k, float tau)
{
	double beta_tol = 1e-6;

	int nch = csize[3],
	    ks1 = ksize[0],
	    ks2 = ksize[1],
	    ks3 = ksize[2];

	int n = ks1 * ks2 * ks3 * nch;

	cplxf *calib;
	int Nc = csize[0]*csize[1]*csize[2]*csize[3];
	cuda (Malloc (&calib, Nc*sizeof(cplxf)));
	cuda (Memcpy (calib, _calib, Nc*sizeof(cplxf), cudaMemcpyHostToDevice));

	cplxf *Q; cuda (Malloc (&Q, n*k*sizeof(cplxf))); //= new cplxf [n*k];
	cplxf *alpha; cuda (Malloc (&alpha, k*sizeof(cplxf)));// = new cplxf [k];
	float *beta; cuda (Malloc (&beta, k*sizeof(float)));// = new float [k];

	cuda (Memset (Q, 0, n*k*sizeof(cplxf)));
	cuda (Memset (alpha, 0, k*sizeof(cplxf)));
	cuda (Memset (beta, 0, k*sizeof(float)));

	cplxf *b = new cplxf[n];
	for (int i = 0; i < n; ++i){
		b[i].real = drand48();
		b[i].imag = drand48();
	}

	double bnrm = 0;
	for (int i = 0; i < n; ++i){ bnrm += abs2 (b[i]); }
	bnrm = sqrt(bnrm);

	for (int i = 0; i < n; ++i){ b[i] *= (1.0 / bnrm); }

	cuda (Memcpy (Q, b, n*sizeof(*Q), cudaMemcpyHostToDevice));

	delete [] b;


	cplxf *z;// = new cplxf[n];
	cplxf *w;// = new cplxf[n];
	cuda (Malloc (&z, n*sizeof(*z)));
	cuda (Malloc (&w, n*sizeof(*w)));

	double ts_iter = timestamp();

	cout << "Lanczos Iterations: " << endl;

	for (int j = 0; j < k; ++j)
	{
		calib3d_cuatamul (z, calib, csize, &Q[j*n], ksize);

		cplxf alpha_j = cublasCdotc (n, (cuComplex*) &Q[j*n], 1,
		                                (cuComplex*) z, 1);
		cuda (Memcpy (&alpha[j], &alpha_j, sizeof(alpha_j),
		               cudaMemcpyHostToDevice));

		for (int t = 0; t < 2; ++t) // reorthogonalize {\emph twice}!
		{
			cuComplex _1 = { 1.0f, 0.0f},
			          _0 = { 0.0f, 0.0f},
			         _n1 = {-1.0f, 0.0f};

			cublasCgemv ('C', n, j+1,
			              _1, (cuComplex*) Q, n, 
			                  (cuComplex*) z, 1,
			              _0, (cuComplex*) w, 1);

			cublasCgemv ('N', n, j+1,
			             _n1, (cuComplex*) Q, n,
			                  (cuComplex*) w, 1,
			             _1,  (cuComplex*) z, 1);
		}

		float znrm = cublasScnrm2 (n, (cuComplex*) z, 1);

		cuda (Memcpy (&beta[j], &znrm, sizeof(znrm), cudaMemcpyHostToDevice));

		if (fabsf (znrm) < beta_tol) break;

		/*
		if (((j+1)%100) == 0){
			double ts = timestamp() - ts_iter;

			cout << "  ... beta[" << j << "] = " << fabsf(znrm)
			     << " (" << ts << " s, " << (ts/k) << " s/iter)" << endl;
		}
		*/

		if (j < k-1){
			cuda (Memcpy (&Q[n*(j+1)], z, n*sizeof(*Q),
			              cudaMemcpyDeviceToDevice));
			cublasCsscal (n, (1.0/znrm), (cuComplex*) &Q[n*(j+1)], 1);
		}
	}

	cuda (ThreadSynchronize());
	cout << "lanczos iteration time: " << (timestamp() - ts_iter) << endl;

	cplxf *_alpha = alpha; alpha = new cplxf [k];
	float *_beta = beta; beta = new float [k];
	cplxf *_Q = Q; Q = new cplxf [n*k];

	cuda (Memcpy (alpha, _alpha, k*sizeof(cplxf), cudaMemcpyDeviceToHost));
	cuda (Memcpy (beta, _beta, k*sizeof(float), cudaMemcpyDeviceToHost));
	cuda (Memcpy (Q, _Q, n*k*sizeof(cplxf), cudaMemcpyDeviceToHost));

	cout << "   beta = " << beta[k-1] << endl;

	cuda (Free (calib));
	cuda (Free (_alpha));
	cuda (Free (_beta));
	cuda (Free (_Q));
	cuda (Free (z));
	cuda (Free (w));

	//debug_shm ("alpha", alpha, 2, k, 1);
	//debug_shm ("beta", beta, 2, k, 1);
	//debug_shm ("Q", Q, 2, n, k);

	float *D = new float [n];
	for (int i = 0; i < n; ++i){ D[i] = alpha[i].real; }
	int info = 0;

	double ts = timestamp();

	csteqr ('V', k, (float*) D, beta, (complex*) Q, n, &info);

	cout << "csteqr time: " << (timestamp() - ts) << endl;

	if (info != 0){
		cerr << "csteqr info=>" << info << endl;
	}

	//debug_shm ("D", D, 2, k, 1);
	//debug_shm ("V", Q, 2, n, k);

	float maxD = fabsf (D[0]);

	for (int i = 1; i < k; ++i)
		if (fabsf (D[i]) > maxD)
			maxD = fabsf (D[i]);

	nv = 0; int *vi = new int[k];
	for (int i = 0; i < k; ++i)
		if (fabsf (D[i]) >= tau*maxD)
			vi[nv++] = i;

	V = new cplxf [nv*n];

	for (int j = 0; j < nv; ++j)
		for (int i = 0; i < n; ++i)
			V[i + n*j] = Q [i + n*vi[j]];

	delete [] alpha;
	delete [] beta;
	delete [] Q;
	delete [] D;
}
